import argparse
import os
import sys
import cv2
import numpy as np
from ultralytics import YOLO
from facenet_pytorch import InceptionResnetV1  # FaceNet模型
from PIL import Image
from torchvision import transforms


def init_args():
    """初始化参数：指定两张对比图片、相似度阈值等"""
    parser = argparse.ArgumentParser(description="YOLOv8+FaceNet 人脸对比（判断两张图是否为同一人）")
    # 必选：两张对比图片的路径（默认当前目录下的 zhou.png 和 zhou1.png）
    parser.add_argument("--img1", type=str,
                        default=os.path.join(os.path.dirname(sys.argv[0]), "zhou.png"),
                        help="第一张图片路径（默认：当前目录/zhou.png）")
    parser.add_argument("--img2", type=str,
                        default=os.path.join(os.path.dirname(sys.argv[0]), "zhou1.png"),
                        help="第二张图片路径（默认：当前目录/zhou1.png）")
    # 可选：相似度阈值（默认0.9，值越高判断越严格，建议0.85-0.95）
    parser.add_argument("--threshold", type=float, default=0.9,
                        help="相似度阈值（>阈值判定为同一人，默认0.9）")
    # 可选：输出目录（保存人脸检测结果图）
    parser.add_argument("--output", type=str, default="./face_comparison_results",
                        help="结果保存目录（默认：./face_comparison_results）")
    return parser.parse_args()


def check_image_file(img_path):
    """检查图片文件是否存在且格式合法"""
    supported_ext = [".jpg", ".jpeg", ".png", ".bmp"]
    if not os.path.exists(img_path):
        raise FileNotFoundError(f"❌ 图片不存在：{img_path}")
    if os.path.splitext(img_path)[1].lower() not in supported_ext:
        raise ValueError(f"❌ 不支持的图片格式：{os.path.splitext(img_path)[1]}（支持：{supported_ext}）")
    # 尝试读取图片，确认文件可正常打开
    img = cv2.imread(img_path)
    if img is None:
        raise RuntimeError(f"❌ 图片损坏或无法读取：{img_path}")
    return img


def load_yolov8_face_model():
    """加载YOLOv8人脸检测模型（yolov8n-face.pt，轻量版适合快速检测）"""
    print("🚀 加载YOLOv8人脸检测模型...")
    try:
        # 自动下载人脸检测模型（约6MB，专为人脸优化）
        model = YOLO("yolov8n-face.pt")
        return model
    except Exception as e:
        raise RuntimeError(
            f"❌ YOLOv8模型加载失败：{str(e)}\n建议手动下载模型到当前目录：https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n-face.pt")


def detect_and_crop_face(yolo_model, img, img_path, output_dir):
    """用YOLOv8检测人脸并裁剪（返回裁剪后的人脸图像，若多个人脸取最大的）"""
    # 执行人脸检测（置信度阈值0.5，过滤低置信度结果）
    results = yolo_model(img, conf=0.5, classes=0, verbose=False)  # classes=0表示只检测人脸
    face_crops = []

    # 提取检测到的人脸坐标并裁剪
    for result in results:
        for box in result.boxes:
            # 获取人脸边界框（xyxy格式：左上x、左上y、右下x、右下y）
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            # 裁剪人脸区域（避免越界）
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(img.shape[1], x2)
            y2 = min(img.shape[0], y2)
            face_crop = img[y1:y2, x1:x2]
            face_crops.append(face_crop)

    # 处理人脸数量异常
    if len(face_crops) == 0:
        raise RuntimeError(f"❌ 未在图片中检测到人脸：{os.path.basename(img_path)}")
    if len(face_crops) > 1:
        print(f"⚠️  图片 {os.path.basename(img_path)} 检测到{len(face_crops)}个人脸，将使用最大的人脸进行对比")
        # 选择最大的人脸（按面积排序）
        face_crops.sort(key=lambda x: x.shape[0] * x.shape[1], reverse=True)

    # 保存带人脸框的原图（方便验证）
    img_with_box = img.copy()
    x1, y1, x2, y2 = map(int, results[0].boxes[0].xyxy[0])  # 最大人脸的框
    cv2.rectangle(img_with_box, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 画绿色框
    cv2.putText(img_with_box, "Face Detected", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    save_path = os.path.join(output_dir, f"detected_{os.path.basename(img_path)}")
    cv2.imwrite(save_path, img_with_box)
    print(f"📁 带人脸框的图片已保存：{save_path}")

    return face_crops[0]  # 返回最大的人脸裁剪图


def load_facenet_model():
    """加载FaceNet模型（用于提取人脸特征向量）"""
    print("\n🚀 加载FaceNet人脸识别模型...")
    try:
        # 加载预训练的FaceNet模型（自动下载约100MB，支持生成128维特征）
        model = InceptionResnetV1(pretrained="vggface2").eval()  # eval()表示推理模式
        # 定义人脸图像预处理（FaceNet要求输入：160x160 RGB图像，归一化）
        preprocess = transforms.Compose([
            transforms.Resize((160, 160)),  # 缩放到160x160
            transforms.ToTensor(),  # 转为Tensor（0-1）
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化到[-1,1]
        ])
        return model, preprocess
    except Exception as e:
        raise RuntimeError(
            f"❌ FaceNet模型加载失败：{str(e)}\n建议检查网络（首次需下载模型）或升级facenet-pytorch：pip install --upgrade facenet-pytorch")


def extract_face_feature(facenet_model, preprocess, face_img):
    """将裁剪后的人脸图像转为128维特征向量"""
    # 转换格式：OpenCV（BGR）→ PIL（RGB）
    face_pil = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))
    # 预处理并转换为特征向量
    img_tensor = preprocess(face_pil).unsqueeze(0)  # 增加batch维度（1,3,160,160）
    with torch.no_grad():  # 关闭梯度计算，加快推理
        feature = facenet_model(img_tensor).numpy()  # 生成128维特征（1,128）
    # 归一化特征向量（提升相似度计算准确性）
    feature = feature / np.linalg.norm(feature, axis=1, keepdims=True)
    return feature[0]  # 返回一维特征向量（128,）


def calculate_cosine_similarity(feat1, feat2):
    """计算两个特征向量的余弦相似度（值范围[-1,1]，越接近1越相似）"""
    similarity = np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2))
    return max(similarity, 0.0)  # 确保相似度非负（极端情况可能为负，取0）


def compare_two_faces(args):
    """完整流程：加载图片→检测人脸→提取特征→计算相似度→判断结果"""
    # 1. 初始化输出目录
    output_dir = os.path.abspath(args.output)
    os.makedirs(output_dir, exist_ok=True)
    print(f"📁 结果保存目录：{output_dir}\n")

    # 2. 检查并加载两张图片
    print("🔍 检查图片文件...")
    img1 = check_image_file(args.img1)
    img2 = check_image_file(args.img2)
    print(f"✅ 图片1：{os.path.basename(args.img1)}（尺寸：{img1.shape[1]}x{img1.shape[0]}）")
    print(f"✅ 图片2：{os.path.basename(args.img2)}（尺寸：{img2.shape[1]}x{img2.shape[0]}）")

    # 3. 用YOLOv8检测并裁剪人脸
    yolo_model = load_yolov8_face_model()
    print("\n👤 检测并裁剪人脸...")
    face1 = detect_and_crop_face(yolo_model, img1, args.img1, output_dir)
    face2 = detect_and_crop_face(yolo_model, img2, args.img2, output_dir)
    print(f"✅ 人脸1裁剪完成（尺寸：{face1.shape[1]}x{face1.shape[0]}）")
    print(f"✅ 人脸2裁剪完成（尺寸：{face2.shape[1]}x{face2.shape[0]}）")

    # 4. 用FaceNet提取人脸特征
    facenet_model, preprocess = load_facenet_model()
    print("\n📊 提取人脸特征向量...")
    feat1 = extract_face_feature(facenet_model, preprocess, face1)
    feat2 = extract_face_feature(facenet_model, preprocess, face2)
    print(f"✅ 人脸1特征：128维向量（示例：{feat1[:5]}...）")
    print(f"✅ 人脸2特征：128维向量（示例：{feat2[:5]}...）")

    # 5. 计算相似度并判断结果
    print("\n⚖️  计算相似度并判断...")
    similarity = calculate_cosine_similarity(feat1, feat2)
    threshold = args.threshold
    is_same_person = similarity >= threshold

    # 6. 输出最终结果
    print("\n" + "=" * 50)
    print("🎯 人脸对比结果")
    print("=" * 50)
    print(f"图片1路径：{args.img1}")
    print(f"图片2路径：{args.img2}")
    print(f"相似度阈值：{threshold}")
    print(f"实际相似度：{similarity:.4f}")
    print(f"结论：{'✅ 两张图片是同一人' if is_same_person else '❌ 两张图片不是同一人'}")
    print("=" * 50)


if __name__ == "__main__":
    # 提前导入依赖（若缺失自动安装）
    required_packages = ["torch", "torchvision", "facenet-pytorch", "opencv-python", "numpy", "pillow"]
    for pkg in required_packages:
        try:
            __import__(pkg)
        except ImportError:
            print(f"❌ 缺少依赖包 {pkg}，正在自动安装...")
            os.system(f"{sys.executable} -m pip install {pkg} --upgrade")

    # 导入PyTorch（FaceNet需要）
    import torch

    try:
        args = init_args()
        compare_two_faces(args)
        print("\n🎉 人脸对比完成！可查看输出目录的检测结果图。")
    except Exception as e:
        print(f"\n❌ 程序执行出错：{str(e)}")
        sys.exit(1)
